#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
视频内容总结生成器
根据已有的编码字段、视频描述和大家都在搜，智能生成视频内容总结
"""

import os
import re
import json
import requests
from typing import Dict, List, Optional
from dotenv import load_dotenv

class VideoSummaryGenerator:
    """
    视频内容总结生成器类
    """
    
    def __init__(self, results_dir: str):
        """
        初始化总结生成器
        
        Args:
            results_dir: 结果文件目录
        """
        self.results_dir = results_dir
        self.deepseek_api_key = os.getenv("DEEPSEEK_API_KEY")
        self.deepseek_api_url = "https://api.deepseek.com/v1/chat/completions"
    
    def _call_deepseek_api(self, prompt: str) -> Optional[str]:
        """
        调用DeepSeek API生成总结
        
        Args:
            prompt: 提示词
            
        Returns:
            生成的总结内容，如果失败返回None
        """
        if not self.deepseek_api_key:
            print("DeepSeek API Key 未设置")
            return None
        
        try:
            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {self.deepseek_api_key}"
            }
            
            payload = {
                "model": "deepseek-chat",
                "messages": [
                    {"role": "system", "content": "你是一个专业的视频内容分析助手，擅长根据视频信息生成准确、简洁的视频内容总结。"},
                    {"role": "user", "content": prompt}
                ],
                "temperature": 0.3,
                "max_tokens": 500
            }
            
            response = requests.post(self.deepseek_api_url, headers=headers, json=payload, timeout=30)
            response.raise_for_status()
            
            result = response.json()
            return result['choices'][0]['message']['content']
        except Exception as e:
            print(f"调用DeepSeek API失败: {e}")
            return None
    
    def _generate_summary_prompt(self, video_desc: str, search_tags: str, encoded_fields: Dict[str, str]) -> str:
        """
        生成用于DeepSeek API的提示词
        
        Args:
            video_desc: 视频描述
            search_tags: 大家都在搜
            encoded_fields: 编码字段字典
            
        Returns:
            格式化的提示词
        """
        encoded_fields_str = "\n".join([f"{key}: {value}" for key, value in encoded_fields.items()])
        
        prompt = f"""
请根据以下视频信息，生成一段100-200字的视频内容总结。
总结应当：
1. 准确描述视频的主要内容和亮点
2. 与编码字段的分类保持一致
3. 为编码字段提供证据支持
4. 语言流畅、自然，符合中文表达习惯
5. 字数必须在100-200字之间
6. 只返回总结内容，不要添加额外的标题或说明

视频信息：
- 视频描述: {video_desc}
- 大家都在搜: {search_tags}
- 编码字段: 
{encoded_fields_str}
"""
        
        return prompt
    
    def _extract_info_from_txt(self, file_path: str) -> Dict[str, any]:
        """
        从TXT文件中提取所需信息
        
        Args:
            file_path: TXT文件路径
            
        Returns:
            包含视频描述、大家都在搜和编码字段的字典
        """
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.readlines()
        
        result = {
            'video_desc': '',
            'search_tags': '',
            'encoded_fields': {},
            'content_lines': content
        }
        
        current_section = None
        
        for line in content:
            stripped_line = line.strip()
            
            # 识别当前部分
            if stripped_line == '=== 视频元数据 ===':
                current_section = 'video_metadata'
                continue
            elif stripped_line == '=== 达人元数据 ===':
                current_section = 'creator_metadata'
                continue
            elif stripped_line == '=== 编码字段 ===':
                current_section = 'encoded_fields'
                continue
            
            # 提取视频描述和大家都在搜
            if current_section == 'video_metadata':
                if stripped_line.startswith('视频描述:'):
                    result['video_desc'] = stripped_line[5:].strip()
                elif stripped_line.startswith('大家都在搜:'):
                    result['search_tags'] = stripped_line[6:].strip()
            
            # 提取编码字段
            elif current_section == 'encoded_fields':
                if ':' in stripped_line:
                    parts = stripped_line.split(':', 1)
                    if len(parts) == 2:
                        field_name = parts[0].strip()
                        field_value = parts[1].strip()
                        result['encoded_fields'][field_name] = field_value
        
        return result
    
    def _update_txt_with_summary(self, file_path: str, info: Dict[str, any], summary: str):
        """
        将总结添加到TXT文件的编码字段上方
        
        Args:
            file_path: TXT文件路径
            info: 提取的信息
            summary: 生成的总结
        """
        content_lines = info['content_lines']
        updated_lines = []
        
        for line in content_lines:
            if line.strip() == '=== 编码字段 ===':
                # 在编码字段标题前添加总结
                updated_lines.append(f"视频内容总结: {summary}\n\n")
            updated_lines.append(line)
        
        # 写入更新后的内容
        with open(file_path, 'w', encoding='utf-8') as f:
            f.writelines(updated_lines)
    
    def process_all_files(self, limit: int = None):
        """
        处理所有TXT文件
        
        Args:
            limit: 处理文件数量限制，None表示处理所有文件
        """
        # 获取所有TXT文件
        txt_files = [f for f in os.listdir(self.results_dir) if f.endswith('.txt')]
        
        # 限制处理文件数量
        if limit is not None:
            txt_files = txt_files[:limit]
        
        processed_count = 0
        success_count = 0
        
        for txt_file in txt_files:
            file_path = os.path.join(self.results_dir, txt_file)
            print(f"处理文件 {processed_count + 1}/{len(txt_files)}: {txt_file}")
            
            try:
                # 提取信息
                info = self._extract_info_from_txt(file_path)
                
                # 检查是否已有总结（避免重复处理）
                has_summary = any(line.strip().startswith('视频内容总结:') for line in info['content_lines'])
                if has_summary:
                    print(f"  文件 {txt_file} 已有总结，跳过处理")
                    processed_count += 1
                    continue
                
                # 生成提示词
                prompt = self._generate_summary_prompt(
                    info['video_desc'], 
                    info['search_tags'], 
                    info['encoded_fields']
                )
                
                # 调用API生成总结
                summary = self._call_deepseek_api(prompt)
                
                if summary:
                    # 确保总结在100-200字之间
                    summary = summary.strip()
                    if len(summary) < 100:
                        print(f"  警告: 生成的总结过短 ({len(summary)}字)")
                    elif len(summary) > 200:
                        print(f"  警告: 生成的总结过长 ({len(summary)}字)，截取到200字")
                        summary = summary[:200]
                    
                    # 更新文件
                    self._update_txt_with_summary(file_path, info, summary)
                    success_count += 1
                    print(f"  成功生成总结")
                else:
                    print(f"  生成总结失败")
                    
            except Exception as e:
                print(f"  处理文件失败: {e}")
            
            processed_count += 1
        
        print(f"\n处理完成!")
        print(f"总共处理: {processed_count} 个文件")
        print(f"成功生成总结: {success_count} 个文件")


def main():
    """
    主函数
    """
    # 加载环境变量
    load_dotenv()
    
    # 设置结果目录
    results_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'results')
    
    # 检查结果目录是否存在
    if not os.path.exists(results_dir):
        print(f"错误: 结果目录 {results_dir} 不存在")
        return
    
    # 创建总结生成器实例
    generator = VideoSummaryGenerator(results_dir)
    
    # 处理所有文件
    generator.process_all_files()


if __name__ == "__main__":
    main()